from typing import List, Dict, Optional
from .ast_nodes import *

class IRGenerator:
    def __init__(self):
        self.lines: List[str] = []
        self.temp_count = 0
        self.label_count = 0
        self.strings: Dict[str, str] = {} # 内容 -> 全局变量名
        self.global_vars: Dict[str, str] = {} # 名称 -> ir_name
        self.local_vars: Dict[str, str] = {} # 名称 -> ir_name (寄存器或alloca)
        self.current_func_ret_type = None
        self.record_defs: Dict[str, RecordDef] = {}
        self.var_types: Dict[str, str] = {}  # Track variable types for FieldAccess

    def generate(self, program: Program) -> str:
        self.emit("; ModuleID = 'ink_module'")
        self.emit("source_filename = \"ink_module\"")
        self.emit("target datalayout = \"e-m:w-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128\"")
        self.emit("target triple = \"x86_64-pc-windows-msvc\"") # 假设 Windows x64
        self.emit("")
        
        # 声明 printf 用于测试
        self.emit("declare i32 @printf(i8*, ...)")
        self.emit("")

        for defn in program.definitions:
            self.visit_definition(defn)
            
        # 发射字符串常量
        for content, name in self.strings.items():
            # 计算包括 null 终止符的长度
            length = len(content.encode('utf-8')) + 1
            # 为 LLVM 转义字符串
            escaped = content.replace('\\', '\\5C').replace('"', '\\22').replace('\n', '\\0A')
            self.emit(f"{name} = private unnamed_addr constant [{length} x i8] c\"{escaped}\\00\", align 1")
            
        return "\n".join(self.lines)

    def quote(self, name: str) -> str:
        """Quote Chinese identifiers for LLVM IR"""
        return f'"{name}"'

    def visit_definition(self, node: ASTNode):
        if isinstance(node, FunctionDef):
            self.visit_function_def(node)
        elif isinstance(node, VariableDef):
            # 全局变量
            pass
        elif isinstance(node, RecordDef):
            # Store definition for later lookup (field indices)
            self.record_defs[node.name] = node
            
            # Define LLVM struct type
            # %"Name" = type { i64, i64, ... }
            field_types = ["i64" for _ in node.fields]
            self.emit(f"%{self.quote(node.name)} = type {{ {', '.join(field_types)} }}")

    def visit_function_def(self, node: FunctionDef):
        self.local_vars = {}
        self.var_types = {}  # Reset variable type tracking
        self.temp_count = 0
        
        # 将墨算类型映射到 LLVM 类型
        ret_type = self.map_type(node.return_type) if node.return_type else "i64" # 默认为 i64
        self.current_func_ret_type = ret_type
        
        params_str = []
        for param in node.params:
            param_type = self.map_type(param.type_annotation) if param.type_annotation else "i64"
            param_name = f"%{param.name}"
            params_str.append(f"{param_type} {param_name}")
            self.local_vars[param.name] = param_name
            
        # 将墨算入口点 "夏" 映射到 LLVM main
        func_name = f"@{self.quote(node.name)}"
        if node.name == "夏":
            func_name = "@main"
             
        self.emit(f"define {ret_type} {func_name}({', '.join(params_str)}) {{")
        self.emit("entry:")
        
        # 生成函数体
        result = self.visit_block(node.body)
        
        # 返回结果
        if ret_type == "void":
            self.emit("  ret void")
        else:
            self.emit(f"  ret {ret_type} {result}")
            
        self.emit("}")
        self.emit("")

    def visit_block(self, node: Block) -> str:
        last_val = "0"
        for stmt in node.statements:
            if isinstance(stmt, Expr):
                last_val = self.visit_expr(stmt)
            elif isinstance(stmt, VariableDef):
                # 本地变量
                val = self.visit_expr(stmt.value)
                self.local_vars[stmt.name] = val
                # Track the type if it's a RecordLiteral
                if isinstance(stmt.value, RecordLiteral):
                    self.var_types[stmt.name] = stmt.value.record_type
                last_val = val
        return last_val

    def visit_expr(self, node: Expr) -> str:
        if isinstance(node, Literal):
            if node.type_name == '数':
                return str(node.value)
            elif node.type_name == '徽数':
                # LLVM 浮点字面量
                return f"0x{float(node.value).hex()}"
            elif node.type_name == '词':
                # 字符串字面量
                # 去掉引号
                content = node.value.strip('"')
                return self.get_string_constant(content)
            elif node.type_name == '两仪':
                # 布尔值
                return "1" if node.value == '真' else "0"
            else:
                return "0"
                
        elif isinstance(node, Identifier):
            return self.local_vars.get(node.name, f"%{node.name}")
            
        elif isinstance(node, BinOp):
            # 简单的二元运算符
            left = self.visit_expr(node.left)
            right = self.visit_expr(node.right)
            
            op_map = {
                '+': 'add',
                '-': 'sub',
                '*': 'mul',
                '/': 'sdiv',
                '%': 'srem',
                '==': 'icmp eq',
                '!=': 'icmp ne',
                '<': 'icmp slt',
                '>': 'icmp sgt',
                '<=': 'icmp sle',
                '>=': 'icmp sge',
            }
            
            op_llvm = op_map.get(node.op, 'add')
            tmp = self.new_temp()
            
            if op_llvm.startswith('icmp'):
                self.emit(f"  {tmp} = {op_llvm} i64 {left}, {right}")
                # 将 i1 转换为 i64
                result = self.new_temp()
                self.emit(f"  {result} = zext i1 {tmp} to i64")
                return result
            else:
                self.emit(f"  {tmp} = {op_llvm} i64 {left}, {right}")
                return tmp
            
        elif isinstance(node, CallExpr):
            # 函数调用
            func = node.func
            args = node.args
            
            # 访问参数
            arg_vals = [f"i64 {self.visit_expr(arg)}" for arg in args]
            
            if isinstance(func, Identifier):
                # 直接函数调用
                func_name = f"@{self.quote(func.name)}"
                if func.name == "夏":
                    func_name = "@main"
                
                # 确定返回类型 - 检查函数是否具有推断类型
                ret_type = "i64"
                if hasattr(node.func, 'inferred_type') and node.func.inferred_type:
                    if isinstance(node.func.inferred_type, FunctionType):
                        ret_type = self.map_type(node.func.inferred_type.return_type)
                
                # 特殊情况：C 函数如 printf
                if func_name == "@printf":
                    ret_type = "i32"
                
                # 特殊情况：书 (print) - 根据类型生成适当的 printf 调用
                if func_name == "@书" or func.name == "书":
                    # 根据参数类型生成适当的 printf 调用
                    if len(args) > 0:
                        arg = args[0]
                        arg_val_with_type = arg_vals[0] # e.g. "i64 %1"
                        # Strip "i64 " prefix to get raw value
                        arg_val = arg_val_with_type.replace("i64 ", "")
                        
                        # 根据类型确定格式字符串
                        if hasattr(arg, 'inferred_type') and arg.inferred_type:
                            arg_type_name = arg.inferred_type.name
                            if arg_type_name == '数':
                                # 打印整数: printf("%lld\n", value)
                                fmt_str = self.get_string_constant("%lld\n")
                                printf_res = self.new_temp()
                                # print(f"DEBUG codegen: arg_val='{arg_val}'")
                                self.emit(f"  {printf_res} = call i32 @printf(i8* {fmt_str}, i64 {arg_val})")
                            elif arg_type_name == '词':
                                # 打印字符串: printf("%s\n", value)
                                fmt_str = self.get_string_constant("%s\n")
                                printf_res = self.new_temp()
                                self.emit(f"  {printf_res} = call i32 @printf(i8* {fmt_str}, i8* {arg_val})")
                            elif arg_type_name == '两仪':
                                # 打印布尔值: 使用 select 选择字符串
                                # 创建 true/false 的字符串常量
                                true_str = self.get_string_constant("真")
                                false_str = self.get_string_constant("假")
                                
                                # 将 i64 转换为 i1 用于 select
                                cond = self.new_temp()
                                self.emit(f"  {cond} = icmp ne i64 {arg_val}, 0")
                                
                                # 选择 true/false 字符串
                                result_str = self.new_temp()
                                self.emit(f"  {result_str} = select i1 {cond}, i8* {true_str}, i8* {false_str}")
                                
                                # 打印选择的字符串
                                fmt_str = self.get_string_constant("%s\n")
                                printf_res = self.new_temp()
                                self.emit(f"  {printf_res} = call i32 @printf(i8* {fmt_str}, i8* {result_str})")
                            else:
                                # 默认：打印为整数
                                fmt_str = self.get_string_constant("%lld\n")
                                printf_res = self.new_temp()
                                self.emit(f"  {printf_res} = call i32 @printf(i8* {fmt_str}, i64 {arg_val})")
                        else:
                            # 没有类型信息，假设为整数
                            fmt_str = self.get_string_constant("%lld\n")
                            printf_res = self.new_temp()
                            self.emit(f"  {printf_res} = call i32 @printf(i8* {fmt_str}, i64 {arg_val})")
                    
                    # 返回 void (用 0 表示)
                    return "0"
                
                tmp = self.new_temp()
                self.emit(f"  {tmp} = call {ret_type} {func_name}({', '.join(arg_vals)})")
                return tmp
            else:
                # 间接调用（函数指针）- 未完全支持
                func_val = self.visit_expr(func)
                tmp = self.new_temp()
                self.emit(f"  {tmp} = call i64 {func_val}({', '.join(arg_vals)})")
                return tmp
            
        elif isinstance(node, IfExpr):
            cond = self.visit_expr(node.condition)
            # Cond is i64 (0 or 1), need i1 for br
            cond_i1 = self.new_temp()
            self.emit(f"  {cond_i1} = icmp ne i64 {cond}, 0")
            
            then_label = self.new_label("then")
            else_label = self.new_label("else")
            merge_label = self.new_label("merge")
            
            self.emit(f"  br i1 {cond_i1}, label %{then_label}, label %{else_label}")
            
            # Then block
            self.emit(f"{then_label}:")
            then_val = self.visit_block(node.then_branch)
            self.emit(f"  br label %{merge_label}")
            
            # Else block
            self.emit(f"{else_label}:")
            else_val = "0"
            if node.else_branch:
                if isinstance(node.else_branch, Block):
                    else_val = self.visit_block(node.else_branch)
                else:
                    else_val = self.visit_expr(node.else_branch)
            self.emit(f"  br label %{merge_label}")
            
            # Merge block
            self.emit(f"{merge_label}:")
            result = self.new_temp()
            self.emit(f"  {result} = phi i64 [{then_val}, %{then_label}], [{else_val}, %{else_label}]")
            return result

        elif isinstance(node, RecordLiteral):
            record_def = self.record_defs.get(node.record_type)
            # Allocate struct
            struct_var = self.new_temp()
            quoted_type = self.quote(node.record_type)
            self.emit(f"  {struct_var} = alloca %{quoted_type}, align 8")
            
            # Store fields
            for i, field in enumerate(record_def.fields):
                # Get value from literal
                val_expr = node.field_values.get(field.name)
                val = self.visit_expr(val_expr)
                
                # GEP and Store
                field_ptr = self.new_temp()
                self.emit(f"  {field_ptr} = getelementptr inbounds %{quoted_type}, %{quoted_type}* {struct_var}, i32 0, i32 {i}")
                self.emit(f"  store i64 {val}, i64* {field_ptr}, align 8")
            
            # Return pointer as i64
            result = self.new_temp()
            self.emit(f"  {result} = ptrtoint %{quoted_type}* {struct_var} to i64")
            return result

        elif isinstance(node, FieldAccess):
            # Get record pointer (as i64)
            record_val = self.visit_expr(node.record)
            
            # Determine record type
            record_type_name = None
            
            # Strategy 1: Use stored record_type_name from sema
            if hasattr(node, 'record_type_name') and node.record_type_name:
                record_type_name = node.record_type_name
            
            # Strategy 2: Check if record is an Identifier and lookup in var_types
            elif isinstance(node.record, Identifier) and node.record.name in self.var_types:
                record_type_name = self.var_types[node.record.name]
            
            # Strategy 3: Use inferred_type if available
            elif hasattr(node.record, 'inferred_type') and node.record.inferred_type:
                if hasattr(node.record.inferred_type, 'name'):
                    record_type_name = node.record.inferred_type.name
            
            if not record_type_name:
                raise RuntimeError(f"Cannot determine record type for field access on '{node.field_name}'")
            
            record_def = self.record_defs.get(record_type_name)
            if not record_def:
                raise RuntimeError(f"Record definition not found for type: {record_type_name}")
            
            quoted_type = self.quote(record_type_name)
            
            # Find field index
            field_index = -1
            for i, f in enumerate(record_def.fields):
                if f.name == node.field_name:
                    field_index = i
                    break
            
            # Cast i64 back to %Name*
            struct_ptr = self.new_temp()
            self.emit(f"  {struct_ptr} = inttoptr i64 {record_val} to %{quoted_type}*")
            
            # GEP and Load
            field_ptr = self.new_temp()
            self.emit(f"  {field_ptr} = getelementptr inbounds %{quoted_type}, %{quoted_type}* {struct_ptr}, i32 0, i32 {field_index}")
            
            result = self.new_temp()
            self.emit(f"  {result} = load i64, i64* {field_ptr}, align 8")
            return result

        return "0"

    def map_type(self, type_node: TypeNode) -> str:
        if not type_node: return "i64"
        if type_node.name == '数': return "i64"  # Use i64 for InkLang integers
        if type_node.name == '两仪': return "i64" # Use i64 for bools
        if type_node.name == '空': return "void"
        return "i64"

    def get_string_constant(self, content: str) -> str:
        if content not in self.strings:
            name = f"@.str.{len(self.strings)}"
            self.strings[content] = name
        
        # 返回第一个字符的指针：getelementptr
        name = self.strings[content]
        length = len(content.encode('utf-8')) + 1
        tmp = self.new_temp()
        self.emit(f"  {tmp} = getelementptr inbounds [{length} x i8], [{length} x i8]* {name}, i64 0, i64 0")
        return tmp

    def new_temp(self) -> str:
        self.temp_count += 1
        return f"%{self.temp_count}"

    def new_label(self, name: str) -> str:
        self.label_count += 1
        return f"{name}.{self.label_count}"

    def emit(self, line: str):
        self.lines.append(line)
